﻿using System;
using System.Collections.Generic;
using System.ComponentModel;
using System.Diagnostics;
using System.IO;
using System.Linq;
using System.Net;
using System.Runtime.InteropServices;
using System.Text;
using Hyperion.Native;

namespace Hyperion.Injector
{
    /// <summary>
    /// Provides support for injecting modules into remote processes.
    /// </summary>
    public class LibraryInjector : IDisposable
    {
        public Dictionary<string, IntPtr> Exports { get; private set; }
        public Dictionary<string, InjectedLibrary> InjectedModules { get; private set; }

        public SafeMemoryHandle Handle { get; private set; }
        public Process Process { get; private set; }
        public ExternalProcessMemory Memory { get; private set; }

        /// <summary>
        /// Gets a value indicating whether we should eject the library from the target process as soon as this injector is disposed.
        /// </summary>
        /// <value>
        ///   <c>true</c> if [eject on dispose]; otherwise, <c>false</c>.
        /// </value>
        public bool EjectOnDispose { get; private set; }

        #region Constructor

        /// <summary>
        /// Initializes a new instance of the <see cref="Injector" /> class.
        /// </summary>
        /// <param name="process">The process.</param>
        /// <param name="ejectOnDispose">if set to <c>true</c> [eject on dispose].</param>
        /// <exception cref="Hyperion.MemoryException">Can not inject stuff into my own process, dumbass.</exception>
        public LibraryInjector(Process process, bool ejectOnDispose)
        {
            if (process.Id == Process.GetCurrentProcess().Id)
                throw new MemoryException("Can not inject stuff into my own process, dumbass.");

            // Let's get us a nice meaty SafeMemoryHandle.
            Handle =
                Win32.OpenProcess(
                    ProcessAccessFlags.QueryInformation | ProcessAccessFlags.CreateThread |
                    ProcessAccessFlags.VMOperation | ProcessAccessFlags.VMRead | ProcessAccessFlags.VMWrite, false,
                    process.Id);

            if (Handle.IsInvalid)
                throw new MemoryException("The target Process Handle is invalid; OpenProcess couldn't be called with the desired access flags. Last Win32 Error:" + Marshal.GetLastWin32Error());

            Process = process;
            InjectedModules = new Dictionary<string, InjectedLibrary>();
            Exports = new Dictionary<string, IntPtr>();
            EjectOnDispose = ejectOnDispose;

            // Get us an ExternalProcessMemory class so we can do some magic like allocating, and stuff.
            Memory = new ExternalProcessMemory(process);
        }

        /// <summary>
        /// Initializes a new instance of the <see cref="Injector"/> class. By default, the injected library is ejected if Injector is initialised this way.
        /// </summary>
        /// <param name="proc">The proc.</param>
        public LibraryInjector(Process proc)
            : this(proc, true)
        {
        }

        #endregion

        #region Injection/Ejection

        /// <summary>
        /// Injects the specified library into the process this injector is initialised around.
        /// </summary>
        /// <param name="libraryPath">The library path.</param>
        /// <returns></returns>
        /// <exception cref="System.IO.FileNotFoundException">Couldn't find the library to inject! Check your stuff!</exception>
        /// <exception cref="Hyperion.MemoryException">
        /// Couldn't inject the specified file because either the Handle is invalid or the Process isn't set!
        /// or
        /// Couldn't obtain a handle for Kernel32! That's bad! [Win32 Error: {0}]
        /// or
        /// Couldn't find the address of LoadLibraryW inside Kernel32! That's bad! [Win32 Error: {0}]
        /// or
        /// Couldn't obtain a pointer to our externally allocated memory! Something went wrong!
        /// or
        /// Couldn't obtain a pointer to the remotely created thread! [Win32 Error: {0}]
        /// or
        /// WaitForSingleObject returned a wrong value after CreateRemoteThread! [Win32 Error: {0}]
        /// or
        /// GetExitCodeThread failed! [Win32 Error: {0}]
        /// or
        /// Couldn't obtain a pointer to the externally injected library, that's bad! Details in inner exception. [Win32 Error: {0}]
        /// or
        /// Couldn't find our injected library in the target process! Injection failed!
        /// </exception>
        /// <exception cref="System.ComponentModel.Win32Exception"></exception>
        public IntPtr Inject(string libraryPath)
        {
            IntPtr libraryPointer = IntPtr.Zero;

            if (!File.Exists(libraryPath))
                throw new FileNotFoundException("Couldn't find the library to inject! Check your stuff!");

            if (Process == null || Handle.IsInvalid)
                throw new MemoryException("Couldn't inject the specified file because either the Handle is invalid or the Process isn't set!");

            var path = Path.GetFullPath(libraryPath);
            var library = Path.GetFileName(path);

            IntPtr threadPtr = IntPtr.Zero;

            try
            {
                // Obtain a handle to kernel32 first, and fail miserably if we can't.
                var kernel32 = Win32.GetModuleHandle("Kernel32");
                if (kernel32 == IntPtr.Zero)
                    throw new MemoryException("Couldn't obtain a handle for Kernel32! That's bad! [Win32 Error: {0}]",
                                              Marshal.GetLastWin32Error());

                // And get a handle for LoadLibraryW while we're at it. Kthx!
                var loadLibraryW = Win32.GetProcAddress(kernel32, "LoadLibraryW");
                if (loadLibraryW == IntPtr.Zero)
                    throw new MemoryException(
                        "Couldn't find the address of LoadLibraryW inside Kernel32! That's bad! [Win32 Error: {0}]",
                        Marshal.GetLastWin32Error());

                // What we'll do is write a pointer to our library's path to an allocated piece of memory, and 
                // we'll then call LoadLibraryW and CreateRemoteThread on it to actually have the target process "spawn" our library
                // and make it alive and doing all sorts of magic.
                byte[] unicodeBytes = Encoding.Unicode.GetBytes(path);

                // Allocate some memory for the full string to our lib and write the pointer to it.
                using (var alloc = Memory.Allocate(unicodeBytes.Length))
                {
                    alloc.WriteBytes(unicodeBytes);

                    if (alloc.Address == IntPtr.Zero)
                        throw new MemoryException("Couldn't obtain a pointer to our externally allocated memory! Something went wrong!");

                    threadPtr = Win32.CreateRemoteThread(Handle.DangerousGetHandle(), IntPtr.Zero, 0x0, loadLibraryW,
                                                             alloc.Address, 0, IntPtr.Zero);

                    if (threadPtr == IntPtr.Zero)
                        throw new MemoryException("Couldn't obtain a pointer to the remotely created thread! [Win32 Error: {0}]", Marshal.GetLastWin32Error());
                }

                if (Win32.WaitForSingleObject(threadPtr, (uint)ThreadWaitValue.Infinite) != (uint)ThreadWaitValue.Object0)
                    throw new MemoryException("WaitForSingleObject returned a wrong value after CreateRemoteThread! [Win32 Error: {0}]", Marshal.GetLastWin32Error());

                var exitCodePtr = IntPtr.Zero;
                if (!Win32.GetExitCodeThread(threadPtr, out exitCodePtr))
                    throw new MemoryException("GetExitCodeThread failed! [Win32 Error: {0}]", Marshal.GetLastWin32Error());

                // If the function succeeds, the return value is nonzero. Meaning that we shouldn't get IntPtr.Zero as long as everything runs properly.
                if (exitCodePtr == IntPtr.Zero)
                    throw new MemoryException("Couldn't obtain a pointer to the externally injected library, that's bad! Details in inner exception. [Win32 Error: {0}]", new Win32Exception(Marshal.GetLastWin32Error()), Marshal.GetLastWin32Error());

                // Let's find out if we actually injected the thing properly, and we can find it.
                ProcessModule module = GetModule(library);

                if (module == null)
                    throw new MemoryException("Couldn't find our injected library in the target process! Injection failed!");

                libraryPointer = module.BaseAddress;
                InjectedModules.Add(library, new InjectedLibrary(library, module));
            }
            finally
            {
                // Clean up that shit please.
                Win32.CloseHandle(threadPtr);
            }

            return libraryPointer;
        }

        /// <summary>
        /// Ejects the specified library from the process we're attached to.
        /// </summary>
        /// <param name="libraryName">Name of the library.</param>
        public void Eject(string libraryName)
        {
            var library = File.Exists(libraryName) ? Path.GetFileName(Path.GetFullPath(libraryName)) : libraryName;

            if (!InjectedModules.ContainsKey(library))
                throw new KeyNotFoundException("The specified module wasn't injected into the target process by this injector, and therefore can't be ejected.");

            var threadPtr = IntPtr.Zero;

            try
            {
                var kernel32 = Win32.GetModuleHandle("Kernel32");
                if (kernel32 == IntPtr.Zero)
                    throw new MemoryException("Couldn't obtain a handle to Kernel32 in the target process! [Win32 Error: {0}]", Marshal.GetLastWin32Error());

                var freeLibraryPtr = Win32.GetProcAddress(kernel32, "FreeLibrary");
                if (freeLibraryPtr == IntPtr.Zero)
                    throw new MemoryException("Couldn't obtain a pointer to FreeLibrary in Kernel32! [Win32 Error: {0}]", Marshal.GetLastWin32Error());

                // Create a thread to call FreeLibrary. Kthx.
                threadPtr = Win32.CreateRemoteThread(Handle.DangerousGetHandle(), IntPtr.Zero, 0x0, freeLibraryPtr,
                                                     InjectedModules[library].BaseAddress, 0, IntPtr.Zero);
                if (threadPtr == IntPtr.Zero)
                    throw new MemoryException("CreateRemoteThread failed in remote process during ejection! [Win32 Error: {0}]", Marshal.GetLastWin32Error());

                if (Win32.WaitForSingleObject(threadPtr, (uint)ThreadWaitValue.Infinite) != (uint)ThreadWaitValue.Object0)
                    throw new MemoryException("WaitForSingleObject didn't return the expected value during eject! [Win32 Error: {0}]", Marshal.GetLastWin32Error());

                var freeLibraryReturnValue = IntPtr.Zero;
                if (!Win32.GetExitCodeThread(threadPtr, out freeLibraryReturnValue))
                    throw new MemoryException("Couldn't obtain the return value of FreeLibrary! [Win32 Error: {0}]", Marshal.GetLastWin32Error());

                if (freeLibraryReturnValue == IntPtr.Zero)
                    throw new MemoryException("FreeLibrary failed in remote thread! The library wasn't successfully ejected from the target process!");
            }
            finally
            {
                Win32.CloseHandle(threadPtr);
            }
        }

        #endregion

        #region Exports

        // TODO: Actually test this stuff properly and add a couple of overloads. Oh, and supporting args properly would be nice too.

        internal IntPtr CallExportInternal(string libraryName, string func, uint timeOut = 30000)
        {
            if (!InjectedModules.ContainsKey(libraryName))
                throw new MemoryException("{0} is not a module injected by this injector!", libraryName);

            var library = File.Exists(libraryName) ? Path.GetFileName(Path.GetFullPath(libraryName)) : libraryName;

            // Obtain a pointer to the export we're after.
            IntPtr exportFuncPtr = InjectedModules[library].ResolveExport(func);
            IntPtr threadPtr = IntPtr.Zero;

            var theirModule = GetModuleBaseAddress(libraryName);
            if (theirModule == IntPtr.Zero)
                throw new MemoryException("Couldn't find the specified module in the remote address; can't call export!");

            // Since the pointer we get from ResolveExport is relative to the module's base, we'll need to rebase it
            // using the base address of "their" version of our module in the process in order to get the final export ptr.
            IntPtr exportPtr = theirModule + (int) exportFuncPtr;
            
            try
            {
                threadPtr = Win32.CreateRemoteThread(Memory.ProcessHandle.DangerousGetHandle(), IntPtr.Zero, 0,
                                                            exportPtr, IntPtr.Zero, 0, IntPtr.Zero);
                if (threadPtr == IntPtr.Zero)
                    throw new MemoryException("CreateRemoteThread failed for calling export {0}!", func, new Win32Exception(Marshal.GetLastWin32Error()));

                uint waitHandle = Win32.WaitForSingleObject(threadPtr, timeOut);
                if (!(waitHandle == (uint) ThreadWaitValue.Object0 || waitHandle == (uint) ThreadWaitValue.Timeout))
                    throw new MemoryException("Waited too long for the thread to exit!", new Win32Exception(Marshal.GetLastWin32Error()));

                IntPtr threadExitCode;
                if (!Win32.GetExitCodeThread(threadPtr, out threadExitCode))
                    throw new MemoryException("Couldn't obtain thread exit code when calling export {0}!", func, new Win32Exception(Marshal.GetLastWin32Error()));

                return threadExitCode;
            }
            finally
            {
                Win32.CloseHandle(threadPtr);
            }
        }

        #endregion

        /// <summary>
        /// Gets the specified module.
        /// </summary>
        /// <param name="moduleName">Name of the module.</param>
        /// <returns></returns>
        public ProcessModule GetModule(string moduleName)
        {
            return Memory.Process.Modules.Cast<ProcessModule>().FirstOrDefault(t => t.ModuleName == moduleName);
        }

        /// <summary>
        /// Gets the base address of a process' module with the specified name.
        /// </summary>
        /// <param name="moduleName">Name of the module.</param>
        /// <returns>The module BaseAddress if successful; otherwise, IntPtr.Zero.</returns>
        public IntPtr GetModuleBaseAddress(string moduleName)
        {
            var module = Memory.Process.Modules.Cast<ProcessModule>().FirstOrDefault(m => m.ModuleName == moduleName);

            return module != null ? module.BaseAddress : IntPtr.Zero;
        }

        #region Implementation of IDisposable

        /// <summary>
        /// Performs application-defined tasks associated with freeing, releasing, or resetting unmanaged resources.
        /// </summary>
        /// <exception cref="System.NotImplementedException"></exception>
        public void Dispose()
        {
            if (Process == null || Handle.IsClosed)
                return; // Nothing to dispose of, kthx.

            if (EjectOnDispose)
            {
                // Get rid of all the modules we injected plx, kthx.
                foreach (var m in InjectedModules)
                    Eject(m.Key);
            }

            Win32.CloseHandle(Handle.DangerousGetHandle());

            // Dispose of the base memory class so it can leave debug mode and clean up and stuff.
            Memory.Dispose();
        }

        #endregion
    }
}
